import argparse
import os
import shutil

import cv2
from imutils import paths
import yaml
from tqdm import tqdm

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="获取杭州人工认为有缺陷大图")
    parser.add_argument('--data',
                        default=r"../data/abnormal.yaml",
                        help="没有分的文件夹")
    parser.add_argument('--saveDir',
                        default=r"C:\Users\zengxh\Desktop\newdatasets",
                        help="没有分的文件夹")
    parser.add_argument('--cleanPathDir',
                        default=r"C:\Users\zengxh\Desktop\solarcell_quejiao", #"../inference/output"
                        help="没有分的文件夹")
    opt = parser.parse_args()

    with open(opt.data,encoding="utf-8") as f:
        data_dict = yaml.load(f, Loader=yaml.FullLoader)  # data dict

    datasetPaths = list(paths.list_images(data_dict["allDatas"]))
    cleanPaths = list(paths.list_images(opt.cleanPathDir))

    datasetNames = [os.path.split(datasetPath)[1] for datasetPath in datasetPaths]

    # 去重
    for cleanPath in tqdm(cleanPaths):
        _, cleanName = os.path.split(cleanPath)
        if cleanName in datasetNames:
            os.remove(os.path.abspath(cleanPath))

    # 人工分类
    cleanPaths = list(paths.list_images(opt.cleanPathDir))
    for cleanPath in tqdm(cleanPaths):
        rgbimg = cv2.imread(cleanPath)
        height, widht, channel = rgbimg.shape

        windowtitle=cleanPath
        # 全屏显示
        # cv2.namedWindow(windowtitle, cv2.WINDOW_NORMAL)  # 改变下参数就可以鼠标随意拖动窗口改变大小
        # cv2.setWindowProperty(windowtitle, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
        # 不是全屏显示
        cv2.namedWindow(windowtitle, cv2.WINDOW_NORMAL)
        cv2.setWindowProperty(windowtitle, cv2.WND_PROP_FULLSCREEN, cv2.WND_PROP_FULLSCREEN)
        cv2.resizeWindow(windowtitle, widht, height)  # 设置显示高宽

        # cv2.moveWindow(windowtitle, 100, 0)  # 设置位置
        cv2.imshow(windowtitle, rgbimg)

        inputval = cv2.waitKey()
        if inputval == 0:  # delete NG
            NGDir = os.path.join(opt.saveDir,"NG")
            if not os.path.exists(NGDir):
                os.makedirs(NGDir)
            shutil.copy(cleanPath,NGDir)
        if inputval == 13:  # 回车 OK
            OKDir = os.path.join(opt.saveDir, "OK")
            if not os.path.exists(OKDir):
                os.makedirs(OKDir)
            shutil.copy(cleanPath, OKDir)
        cv2.destroyAllWindows()